#!/usr/bin/env python

import yaml
import rospy
import rospkg
from geometry_msgs.msg import TransformStamped
from interbotix_common_modules.angle_manipulation import quaternion_is_valid
from interbotix_landmark_modules.landmark import Landmark, LandmarkCollection 
from interbotix_perception_modules.apriltag import InterbotixAprilTagInterface as AprilTag

class LandmarkFinder(object):
    """Application to find landmarks during map exploration"""
    def __init__(self):
        apriltag_ns = rospy.get_param("~apriltag_ns", "apriltag")
        landmark_ns = rospy.get_param("landmark_ns", "landmarks")

        r = rospkg.RosPack()
        pkg_path = r.get_path("interbotix_landmark_modules")

        # get parameters
        finder_ns = rospy.get_namespace().strip("/")
        self.lm_config_filepath = rospy.get_param(
            "~landmark_config",
            "{}/landmarks/landmarks.yaml".format(pkg_path))
        self.obs_frame = rospy.get_param(
            "~obs_frame",
            "camera_color_optical_frame")
        self.fixed_frame = rospy.get_param(
            "~fixed_frame",
            "landmarks")

        self.static_tf_pub = rospy.Publisher(
            'static_transforms', 
            TransformStamped,
            queue_size=10,
            latch=True)

        self.valid_tags = None
        self.landmarks = LandmarkCollection(
            landmarks={},
            obs_frame=self.obs_frame,
            fixed_frame=self.fixed_frame,
            ros_on=True)
        
        self._load_landmarks()
        self.apriltag = AprilTag(
            apriltag_ns=apriltag_ns,
            init_node=False)
        self.apriltag.set_valid_tags(self.valid_tags)
        
        # wait a bit for other nodes to launch
        rospy.sleep(rospy.Duration(5))

        self.timer = rospy.Timer(
            period=rospy.Duration(2),
            callback=self.timer_callback)
        rospy.loginfo("Initialized Landmark Finder")

        rospy.on_shutdown(self._save_landmarks)

    def _load_landmarks(self):
        """loads collection of landmarks from the specified filepath"""
        # load landmarks from filepath
        if self.landmarks.load(self.lm_config_filepath):
            # set tags seen in landmark file
            self.valid_tags = [l.get_id() for l in self.landmarks.data.values()]
            rospy.loginfo("Loaded landmarks from " + self.lm_config_filepath)

    def _save_landmarks(self, tags=None):
        """saves collection of landmarks to the specified filepath"""
        self.landmarks.save(self.lm_config_filepath, tags)

    def timer_callback(self, timer):
        """timed callback that gets the landmark's pose if it is view of the
            camera, transforms the pose to the world's fixed frame, and 
            publishes the TF to the static transform publisher
        """
        # get pose with respect to camera frame
        poses_wrt_cam, tag_ids = self.apriltag.find_pose_id()
        # loop through detected tags, if none, skip
        num_tags = len(tag_ids)
        if num_tags > 0:
            for i in range(num_tags):
                # if we get a valid pose (if given quaternion is valid)...
                if quaternion_is_valid(poses_wrt_cam[i].orientation):
                    # transform the pose to the new frame
                    self.landmarks.data[tag_ids[i]].set_tf_wrt_cam(poses_wrt_cam[i])
                    self.landmarks.data[tag_ids[i]].update_tfs(
                        parent_old=self.obs_frame,
                        parent_new=self.fixed_frame)
            # publish
            self.landmarks.pub_tfs(tag_ids)
            rospy.sleep(rospy.Duration(1.0)) # wait to make sure tf is published
            self.landmarks.pub_markers(tag_ids)


if __name__ == "__main__":
    rospy.init_node("landmark_finder")
    landmark_finder = LandmarkFinder()
    rospy.spin()
